"""
Time Data module. Classes to represent some types of time data.

Parts from pyalgotrade.dataseries [1] used

References:
[1] PyAlgoTrade package: http://gbeced.github.io/pyalgotrade/
"""
from numbers import Number
import math
import numpy as np
import numpy.lib.recfunctions as rf
import datetime
import time
import logging
import misc
import quoteperiod as qp

# TimeData slicing statuses
NOTNEEDED = 0 # Not needed because requested range encompasses existing range
              # completely
EMPTY = 1 # TimeData object is emptgy
DOSLICE = 2 # Slicing must occur
OUTOFRANGE = 3 # Requested range is completely out of existing range

def resize_list(list_, size):
    return list_[-1*size:]

class QuoteStatus:
    """
    Class to concentrate the possible values of a quote status.
    
    I decided not to use enum() so I could document the possible statuses.
    """
    Nothing = 0
    Regular = 1
    AfterMarket = 2
    Auction = 3


class TimeData(object):
    """Time Data base class."""
    
    timestamp = np.array([], dtype=int)
    # This "timestamp2" field is a warped version of timestamp, calculated by
    # misc.TrimGaps(). The field is optional
    timestamp2 = None #np.array([], dtype=float)
    
    # Please fill this in
    attrlist = []
    # Please set this. This is one of the IsDescription classes from within the
    # database module
    _s_dbDescr = None
    
    def __init__(self, *args, **kwargs):
        # =None. maxLen property. None means that this is not controlled
        self.__maxLen = None
    
    def __len__(self):
        """Number of data points based on self.timestamp."""
        return len(self.timestamp)
        
    def __getitem__(self, slice_):
        if slice_.__class__ != slice:
            slice_ = slice(slice_, slice_+1, None)
        list = self.attrlist
        Q = self.__class__()
        for s in self.attrlist:
            Q.__setattr__(s, self.__getattribute__(s)[slice_.start:slice_.stop:slice_.step])
        if self.timestamp2 is not None:
            Q.timestamp2 = self.timestamp2[slice_.start:slice_.stop:slice_.step]
        return Q
    
    def SetSize(self, iSize):
        self.timestamp = np.resize(self.timestamp, iSize)
        # timestamp2 is ignored
        
    def Concatenate(self, a):
        """Concatenates several objects into one.
        
        All elements in a should "ideally" be of same class as self."""
        q = self.__class__()
        if len(a) > 0:
            for s in self.attrlist:
                q.__setattr__(s, np.concatenate([x.__getattribute__(s) for x in a]))
        return q

    def SetSize(self, iSize):
        for s in self.attrlist:
            self.__setattr__(s, np.resize(self.__getattribute__(s), iSize))
        
    def SetData(self, A):
        """Sets fields with recArray from database"""
        
        for s in self.attrlist:
            self.__setattr__(s, A[s])

    def Pack(self):
        """Packs a elements in attrlist using numpy's merge_arrays."""
###        logging.debug("WWWWWWWWWWWWWWWWWW %s %s" % (self.attrlist, self.__class__))
        a = zip(*[self.__getattribute__(s) for s in self.attrlist])
        return a
        #return np.transpose(a)
        
    def FlatCopy(self, class_=None):
        """Creates new TimeData of class_ with same timestamps of self.
        
        Other fields will be vectors with zeros (hence "flat")."""
        r = self.__class__() if class_ is None else class_()
        r.SetSize(len(self))
        r.timestamp = np.copy(self.timestamp)
        if not misc.isempty(self.timestamp):
            r.timestamp2 = np.copy(self.timestamp2)
        return r

    def EmptyCopy(self):
        """Creates new with zero size."""
        r = self.__class__() ###if class_ is None else class_()
        return r

    def GetGeneric(self, fieldName, flagCopy=True):
        """Returns a timedata.Generic containing only one of the y-fields.
        
        For example, calling GetGeneric("close") will return a timedata.Generic
        that will have:
            - timestamp: same as self
            - y: self.close
            
        Arguments:
            fieldName -- must be in self.attrlist
            flagCopy=True -- whether to copy the vectors. Set to False if you
                             don't need to change the values of the result,
                             as it will be faster.
        """
        r = Generic()
        if flagCopy:
            r.timestamp = np.copy(self.timestamp)
            r.y = np.copy(self.__getattribute__(fieldName))
        else:
            r.timestamp = self.timestamp
            r.y = self.__getattribute__(fieldName)
        return r
    
    
    ## Implementation of a "Data Window".
    ## This feature allows to maintain data fields of fixed size. When new
    ## data is added with Append(), old data (from the beginning of the data
    ## vectors) is discarded to have all data vectors of size <= maxLen

    @property
    def maxLen(self):
        return self.__maxLen
    @maxLen.setter
    def maxLen(self, x):
        self.__maxLen = x

        self._CorrectLen()

    def Append(self, td):
        """Appends values at the end, respecting maxLen.Will eventually pop elements from the beginning.
        """
        for s in self.attrlist:
            self.__setattr__(s, np.concatenate(td.__getattribute__(s)))
        self._CorrectLen()

    def _CorrectLen(self):
        """Removes elements from the beginning if len(self) > self.maxLen ."""
        if len(self) > self.__maxLen:
            for s in self.attrlist:
                self.__setattr__(s, resize_list(self.__getattribute__(s),
                                                self.__maxLen))


    
    
    ## SLICING
    ## This feature allows to get slices passing boundary timestamps.
    ## It uses binary search so that
    ## effective_ts1 >= ts1 and effective_ts2 <= ts2


    def GetTimestampSlice(self, ts1, ts2):
        """Returns sliced object of same class."""
        (status, (idx1, idx2)) = self.GetTimestampSliceIndexes(ts1, ts2)
        logging.debug("Status=%d; idx1=%s; idx2=%s" % (status, idx1, idx2))
        return self if status in (NOTNEEDED, EMPTY) \
               else self[idx1:idx2] if status == DOSLICE  \
               else self.EmptyCopy() if status == OUTOFRANGE \
               else self
               
            
    def GetTimestampSliceIndexes(self, ts1, ts2):
        """Returns indexes for TimestampSlice().
        
        Returns:
           (status, (idx1, idx2)), check statuses at the beginning of this file.
           (idx1, idx2) are numbers when status is DOSLICE or NOTNEEDED, otherwise (None, None)
        
        NOTE: status NOTNEEDED returns indexes (0, len(self))
        """
        if len(self) == 0:
            return (EMPTY, (None, None))
###        nan = float('nan')
        xts1 = self.timestamp[0]
        xts2 = self.timestamp[-1]
        status = NOTNEEDED
        if ts1 > xts1:
            if ts1 > xts2:
                logging.debug("ts1 = %d; xts2 = %d" % (ts1, xts2))
                return (OUTOFRANGE, (None, None))
            status = DOSLICE
            idx1 = misc.BSearchCeil(self.timestamp, ts1)
        else:
            idx1 = 0
        if ts2 < xts2:
            if ts2 < xts1:
                logging.debug("ts2 = %d; xts1 = %d" % (ts2, xts1))
                return (OUTOFRANGE, (None, None))
            status = DOSLICE
            idx2 = misc.BSearchFloor(self.timestamp, ts2)+1
        else:
            idx2 = len(self)
        
        return (status, (idx1, idx2))


    def LaggedFlatCopy(self, lag, newLen=None):
        """Creates new TimeData that is smaller than the input.
        
        Arguments:
            lag -- Integer and > 0 specifying number of time shifts into the
                   future.
            newLen -- Size of new TimeData
        
        timestamp comes sized and filled. Uses self["input"][0] to get
        timestamp[lag:]
        
        If newLen is not passed, it will be automatically calculated as the
        original size minus lag.
        
        If the new timestamp vector extends into the future beyond the existing
        timestamp vector, it will be made up somehow
        ."""
        len0 = len(self)
        if newLen is None:
            newLen = max(len0-lag, 0)
        r = self.__class__()
        if newLen > 0:
            r.SetSize(newLen)
        else:
            return r
        if lag < len0:
            idx1, idx2 = lag, min(lag+newLen, len0)
            t = np.copy(self.timestamp[idx1:idx2])
            r.timestamp[0:len(t)] = t
        if newLen+lag > len0:
            dif = newLen+lag-len0
            idxs = np.arange(max(len0, lag), newLen+lag) # Indexes to be used as interpolation input

            print "qqqqq", self.timestamp
            t1 = self.timestamp[-1]
            t0 = self.timestamp[0]
            tspan = float(t1-t0)*(len0+1)/len0
            print "dif", dif, "t1", t1, "t0", t0, "len0", len0
            print "idxs", idxs
            print "coiso", tspan/len0*(idxs-len0)
#            r.timestamp[-dif:] = t1+np.round(tspan/len0*(idxs-len0)) # line fit + extrapolation
            r.timestamp[-dif:] = t0+np.round(tspan/len0*(idxs)) # line fit + extrapolation
        return r
        
    def GetTextTable(self, i0=0, i1=None):
        """Generates a table in text format."""
        fsum = lambda x, y: x+y
        if i1 is None:
            i1 = len(self)
        a = [] # Will concatenate later
        a.append("timestamp           "+reduce(fsum, 
            ["%13s" % x for x in self.attrlist[1:]])) # header
        a.append("-"*len(a[0]))
        for i in range(i0, i1):
            a.append(("%19s " % misc.ts2str(self.timestamp[i]))+
                reduce(fsum, ["%13g" % self.__getattribute__(x)[i] for x in self.attrlist[1:]]))
        return reduce(lambda x, y: x+"\n"+y, a)

    


class Generic(TimeData):
    """Represents single time series (timestamp, y-value)."""
    attrlist = ["timestamp", "y"]
    y = np.array([], dtype=float)
    
    # This cannot be the class itself because it cannot be pickled
    _s_dbDescr = "GenericTimeData"


class Quotes(TimeData):
    """Stores data for open, close, high, low, volume, numTrades"""
    
    attrlist = ["timestamp", "open", "close", "high", "low", "volume",
                    "numTrades", "status", "flagRealtime"]
    
    open = np.array([], dtype=float)
    close = np.array([], dtype=float)
    high = np.array([], dtype=float)
    low = np.array([], dtype=float)
    volume = np.array([], dtype=float)
    numTrades = np.array([], dtype=int)
    status = np.array([], dtype=int)
    flagRealtime = np.array([], dtype=bool)

    _s_dbDescr = "Quote"

################################################################################


def GroupQuotes(q, nm):
    """1-min to 3-min, 5-min etc
    
        q -- timedata.Quotes object
        nm -- number of minutes (must be a divider of 60)
    """
    
    '''
    Should I find something and try to figure out the next boundary, or
    ...'''
    r = Quotes()
    n = len(q)
    if n == 0:
        return r
    incr = nm*60 # Increment in seconds
    r.SetSize(n) # Theoretically, result could be as large as original
                 # (e.g. there is only one quote every nm minutes)
    tt = np.copy(q.timestamp)
    dt0 = misc.ts2dt(tt[0])
    # Finds minimum timestamp that would be included in first block.
    # This will be later used to determine beginnings of other blocks
    tCanonical = misc.dt2ts(dt0.replace(minute=(dt0.minute)//nm*nm))
    i = 0 # Points to original timestamp
    j = 0 # Points to result
    while i < n:
        if i == 0:
            tCanNow = tCanonical # Minimum value in current block
        else:
            if tt[i] < tCanNext+incr: # i.e., no gaps (i.e., empty blocks)
                tCanNow = tCanNext
            else:
                tCanNow = tt[i]-(tt[i]-tCanonical)%(nm*60)
        tCanNext = tCanNow+incr # Minimum value to start next block (at least)
        i0 = i # Points to beginning of block
        while i < n and tt[i] < tCanNext:
            i += 1
        
        # "Recording"
        r.timestamp[j] = tCanNext-60 # 1 minute before beginning of next block #tt[i0]
        r.open[j] = q.open[i0]
        r.close[j] = q.close[i-1]
        r.high[j] = np.max(q.high[i0:i])
        r.low[j] = np.min(q.low[i0:i])
        r.volume[j] = np.sum(q.volume[i0:i])
        r.numTrades[j] = np.sum(q.numTrades[i0:i])
        r.status[j] = q.status[i0]
        r.flagRealtime[j] = q.flagRealtime[i0]
        j += 1
        
###        print "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
###        from PyQt4.QtCore import pyqtRemoveInputHook
###        from pdb import set_trace
###        pyqtRemoveInputHook()
###        set_trace()
        
    r.SetSize(j)
        
    return r
    
    #identifies very beginning
    

def TrimTimeGaps(dd, period):
    """Finds common time gaps across a list of TimeData objects and deletes them.
    
    Returns: modified dd
    
    """

    # Copies timestamp to timestamp2
    for d in dd:
        d.timestamp2 = np.copy(d.timestamp)
        
    if period >= qp.daily:
        return dd # Does nothing if not intraday


    
    #gr = GapRemaker()
    n = len(dd)
    ptrs = np.zeros(n, dtype=int)
    lens = [len(d) for d in dd]
    temp = np.zeros(n)
    breakss = [[] for whocares in range(n)] # List of lists of pointers to shift down everything later on
    oneDay = datetime.timedelta(days=1)
    bFirst = True
    addFactor = 0
    numSec = qp.data[period]["nm"]*60
    prevMaxTime = -numSec # numbers will be first pushed to start at zero with this
    while True:
        # Gets minimum time of the day for each series
        mins = []
        for (d, ptr, len_) in zip(dd, ptrs, lens):
            if ptr < len_:
                mins.append(d.timestamp[ptr])
                
        if len(mins) == 0:
            break # This is the exit door

        currMinTime = np.min(mins)
        gapSize = currMinTime-prevMaxTime-numSec
        addFactor += gapSize # cumulative add factor
        #gr.Add(currMinTime-addFactor, addFactor)
        
        for (ptr, len_, breaks) in zip(ptrs, lens, breakss):
            if len(breaks) == 0 or breaks[-1][0] != ptr:
                breaks.append([ptr, addFactor])
            else:
                breaks[-1][1] += gapSize # It may happen that one pointer stays the same for more than 1 iterations
                                         # This is the case when there are days missing in one of the series
                
        nextDay = misc.dt2ts(misc.ts2dt(currMinTime).replace(hour=0, minute=0)+
                                          oneDay)
        maxs = []
        for (i, d) in enumerate(dd):
            if ptrs[i] <= lens[i]:
                # Increments pointers until day changes or reaches end
                while ptrs[i] < lens[i] and d.timestamp[ptrs[i]] < nextDay:
                    ptrs[i] += 1
                if ptrs[i] > 0:
                    maxs.append(d.timestamp[ptrs[i]-1]) # Maximum time of the day for each series
        if len(maxs) == 0:
            break # OK, another exit door
        prevMaxTime = np.max(maxs)
        bFirst = False

    for (d, breaks) in zip(dd, breakss):
        lb = len(breaks)
        for i in range(len(breaks)):
            (ptr, gapSize) = breaks[i]
            if i == lb-1:
                d.timestamp2[ptr:] -= gapSize
            else:
                d.timestamp2[ptr:breaks[i+1][0]] -= gapSize
    
    return dd


def IntersectTimestamps(dd):
    """Intersection of timestamp (in sync with timestamp2) of QuoteData list.
   
    Returns (timestamp, timestamp2) containing the maximum number of
    elements possible."""
    ii = [i[0] for i in sorted(enumerate(dd), key=lambda x:len(x[1]), reverse=True)]
    
    for i in range(len(ii)):
        if i == 0:
            # Starts composing results
            tto = dd[ii[i]].timestamp # Original timestamps
            ttw = dd[ii[i]].timestamp2 # Warped timestamps
        else:
            if len(dd[ii[i]]) == 0:
                continue # Skips zero-sized data, if any
            to = dd[ii[i]].timestamp
            tw = dd[ii[i]].timestamp2
            if to[0] < tto[0]: # Has something to add at the beginning
                iCut = misc.BSearchRound(to, tto[0]) # 
                if to[iCut] > tto[0]: iCut -= 1
                tto = np.hstack((to[0:iCut+1], tto)) # Prepends
                ttw = np.hstack((tw[0:iCut+1], ttw)) # Mirror operation
            if to[-1] > tto[-1]: # Has something to append at the end
                iCut = misc.BSearchRound(to, tto[-1])
                if to[iCut] < tto[-1]: iCut += 1
                tto = np.hstack((tto, to[iCut:]))
                ttw = np.hstack((ttw, tw[iCut:]))
#                tto.extend(to[iCut:])
#                ttw.extend(tw[iCut:])
                    
        
    return (tto, ttw)
